#!/usr/bin/python

import sys
import re
import numpy as np
#from numpy import *

np.set_printoptions(threshold='nan')

ptt_files_list = sys.argv[1]
cog_id_list = sys.argv[2]

pf = open(ptt_files_list, "rU")
ptt_lines = pf.readlines()

cf = open(cog_id_list, "rU")
cog_lines = cf.readlines()

outfile = "test.out.matrix"

cog_ids = []
for line in cog_lines:
    cog_ids.append(line.rstrip())

ptt_files = []
for line in ptt_lines:
    ptt_files.append(line.rstrip())

x = len(cog_ids)
y = len(ptt_files)

final_matrix = np.zeros([x,y])

#print final_matrix.ndim, final_matrix.shape

header_list = []
m1 = re.compile('(.*),.*')
m2 = re.compile('\d+\.\.\d+.*')

org_names = []
for index, line in enumerate(ptt_files):
    pfh = open(line, "rU")
    ptt = pfh.readlines()
    header = ptt[0].rstrip()
    ptt_cogs_list = []
    if m1.match(header):
        pattern = m1.match(header)
        org_names.append(pattern.group(1))
    for pl in ptt:
        if m2.match(pl):
            pattern = m2.match(pl)
            l = pl.split('\t')
            ptt_cog = l[7]
            ptt_cog_id = ptt_cog[:7]
            #print cog_id
            ptt_cogs_list.append(ptt_cog_id)
        else:
            pass
    for i, cog in enumerate(cog_ids):
        final_matrix[i,index] = ptt_cogs_list.count(cog)
        #print "Done with (" + str(index) + "," + str(i) + ")"

#for org in org_names:
    #print org + "\t",

for row in final_matrix:
    print row

for i, cog in enumerate(cog_ids):
    for j, count in enumerate(ptt_files):
        print final_matrix[i, j]

pf.close()
cf.close()